#include "../DWMRICore/Scalar.h"

#include <stdlib.h>
#include <stdio.h>
#include <math.h>

char *pFAPathname = NULL;			/* input fa file pathname */
char *pBrainMaskPathname = NULL;	/* input brain mask file pathname */
char *pWMMaskPathname = NULL;		/* output white matter mask file pathname */
float threshold = 0.15f;				/* fa value threshold for white matter */

void PrintUsage()
{
	printf("Usage: CreateWhiteMatterMask -fa <fa_file> [-mask <brain_mask_file>] -dst <dst_file>\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fa") == 0) {
			pFAPathname = argv[++i];
		} else if (strcmp(argv[i], "-mask") == 0) {
			pBrainMaskPathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			pWMMaskPathname = argv[++i];
		} else {
			return -1;
		}
	}

	if (pFAPathname == NULL || pWMMaskPathname == NULL)
		return -2;

	return 0;
}

float UpdateP(const uint x, const uint y, const uint z, 
	const uint w, const uint h, const uint d, float *p, float *np)
{
	uint min_x = x > 0 ? x-1 : 0;
	uint max_x = x < w-1 ? x+1 : w-1;
	uint min_y = y > 0 ? y-1 : 0;
	uint max_y = y < h-1 ? y+1 : h-1;
	uint min_z = z > 0 ? z-1 : 0;
	uint max_z = z < d-1 ? z+1 : d-1;

	float dp = 0.0f;
	float dnp = 0.0f;
	uint index = 0;
	for (uint cz = min_z; cz <= max_z; ++cz) {
		for (uint cy = min_y; cy <= max_y; ++cy) {
			for (uint cx = min_x; cx <= max_x; ++cx) {
				index = (cz * h + cy) * w + cx;
				dp += p[index];
				dnp += np[index];
			}
		}
	}

	index = (z * h + y) * w + x;
	float res = p[index] * (1.0f + dp) / (p[index] * (1.0f + dp) + np[index] * (1.0f + dnp));
	return res;
}

void CreateWhiteMatterMask()
{
	CScalar *pFA = new CScalar();
	pFA->OpenScalarFile(pFAPathname);

	CScalar *pBrainMask = NULL;
	if (pBrainMaskPathname != NULL) {
		pBrainMask = new CScalar();
		pBrainMask->OpenScalarFile(pBrainMaskPathname);
	}

	uint w = pFA->GetWidth();
	uint h = pFA->GetHeight();
	uint d = pFA->GetDepth();

	float *p = new float[w*h*d];
	float *np = new float[w*h*d];
	for (uint i = 0; i < w * h * d; ++i) {
		if (pBrainMask == NULL) {
			float fa = pFA->GetScalarFast(i);
			p[i] = 0.5f + (fa - threshold) / (1.0f - threshold) * 0.5f;
			np[i] = 1.0f - p[i];
		} else {
			if (pBrainMask->GetScalarFast(i) > 0.5f) {
				float fa = pFA->GetScalarFast(i);
				p[i] = 0.5f + (fa - threshold) / (1.0f - threshold) * 0.5f;
				np[i] = 1.0f - p[i];
			} else {
				p[i] = 0.0f;
				np[i] = 1.0f;
			}
		}
		
	}

	float *temp_p = new float[w*h*d];
	float *temp_np = new float[w*h*d];
	while (true) {
		for (uint z = 0; z < d; ++z) {
			for (uint y = 0; y < h; ++y) {
				for (uint x = 0; x < w; ++x) {
					uint index = (z * h + y) * w + x;
					temp_p[index] = UpdateP(x, y, z, w, h, d, p, np);
					temp_np[index] = 1.0f - temp_p[index];
				}
			}
		}

		float average_delta = 0.0f;
		for (uint i = 0; i < h*w*d; ++i) {
			float delta = fabs(p[i] - temp_p[i]);
			average_delta += delta;
		}
		average_delta = average_delta / (float)(w*h*d);
		printf("average delta %f\n", average_delta);

		memcpy(p, temp_p, sizeof(float)*w*h*d);
		memcpy(np, temp_np, sizeof(float)*w*h*d);

		if (average_delta < 0.01f)
			break;
	}

	CScalar *pResult = new CScalar();
	pResult->CreateScalar(w, h, d);
	for (uint i = 0; i < w*h*d; ++i) {
		if (p[i] > 0.5f) {
			pResult->SetScalar(i, 1.0f);
		} else {
			pResult->SetScalar(i, 0.0f);
		}
	}
	pResult->SaveScalarFile(pWMMaskPathname);

	SafeDelete(pFA);
	SafeDelete(pResult);
	SafeDeleteArray(p);
	SafeDeleteArray(np);
	SafeDeleteArray(temp_p);
	SafeDeleteArray(temp_np);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	CreateWhiteMatterMask();

	return 0;
}
